#include "OrthographicCamera.cuh"

using namespace OBR;

extern "C"
{
    __constant__ Params params;
}

const float jitter_scale = 1.0;

// #define SSAA
const float ssaa_sample = 55;

extern "C" __global__ void __raygen__orthographic_camera()
{
    RaygenData *sbt_record = reinterpret_cast<RaygenData *>(optixGetSbtDataPointer());
    OrthographicCameraData *camera_data =
        reinterpret_cast<OrthographicCameraData *>(sbt_record->camera_payload);
    const uint3 launch_index = optixGetLaunchIndex();
    const uint3 launch_dims = optixGetLaunchDimensions();
    const uint2 pos = make_uint2(launch_index.x, launch_index.y);
    const unsigned int image_index = launch_index.y * params.width + launch_index.x;
    bool log = SHADER_LOG_COND;

#ifdef DEBUG_SHADER
    if (log)
    {
        printf("Camera data:\n");
        printf("orig: ");
        print_float3(camera_data->orig);
        printf("\n");

        printf("U: ");
        print_float3(camera_data->U);
        printf("\n");

        printf("V: ");
        print_float3(camera_data->V);
        printf("\n");

        printf("dir: ");
        print_float3(camera_data->dir);
        printf("\n");
    }
#endif

    // init RNG
    curandState_t rng;
    curand_init(launch_index.x, launch_index.y, params.launch_time, &rng);
    curandStatePhilox4_32_10_t rng_long;
    curand_init(launch_index.y, launch_index.x, params.launch_time, &rng_long);

    // init prd
    SamplePRD prd;
    prd.params = &params;
    prd.launch_index = launch_index;
    prd.log = log;
    prd.rng = &rng;
    prd.rng_long = &rng_long;

    Spectrum final_radiance;

#ifdef SSAA
    for (size_t aas = 0; aas < ssaa_sample; aas++)
    {
#endif
        for (size_t s = 0; s < params.samples_per_launch; s++)
        {
#ifdef DEBUG_SHADER
            if (log)
            {
                printf("sample %d\n", s);
            }
#endif

            // generate ray
#ifdef SSAA
            float2 subpixel_jitter = Fibonacci2d(aas, ssaa_sample) - 0.5;
#else
        float2 subpixel_jitter = make_float2(curand_uniform(&rng), curand_uniform(&rng)) - 0.5;
        subpixel_jitter *= jitter_scale;
#endif
            float2 inv_filmsize = 1.0f / make_float2(params.width, params.height);
            float2 d = make_float2(static_cast<float>(launch_index.x) + 0.5f + subpixel_jitter.x,
                                   static_cast<float>(launch_index.y) + 0.5f + subpixel_jitter.y);
            RayData ray;
            ray.o = (d.x * inv_filmsize.x - 0.5f) * camera_data->U +
                    (d.y * inv_filmsize.y - 0.5f) * camera_data->V + camera_data->orig;
            ray.d = camera_data->dir;
            ray.med_id = sbt_record->med_id;
            ray.med_instance_id = sbt_record->med_instance_id;
            ray.has_differential = true;
            ray.rx_o = ray.o + inv_filmsize.x * camera_data->U;
            ray.rx_d = camera_data->dir;
            ray.ry_o = ray.o + inv_filmsize.y * camera_data->V;
            ray.ry_d = camera_data->dir;

#ifdef DEBUG_SHADER
            if (log)
            {
                printf("generated ray orig(%f, %f, %f), dir(%f, %f, %f)\n", ray.o.x, ray.o.y,
                       ray.o.z, ray.d.x, ray.d.y, ray.d.z);
            }
#endif

            // set sample specified params
            prd.scene_time = curand_uniform(&rng) * (params.tmax - params.tmin) + params.tmin;
            prd.ray_occlusion_test = false;
            prd.inter.wi = ray;
            // call integrator
            final_radiance += optixContinuationCall<Spectrum, SamplePRD *>(0, &prd);
        }
#ifdef SSAA
    }
    final_radiance /= ssaa_sample;
#endif

    // set radiance to fbo #0
    set_fbo(&params, launch_index, 0, final_radiance / params.samples_per_launch);
}